#!/usr/bin/env python3
# coding=utf-8

#
# Copyright (c) 2022 Huawei Device Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import json
import traceback
import time
import re
import zipfile

from xdevice import AgentMode
from xdevice import Variables
from xdevice import check_uitest_version
from devicetest import RESOURCE_PATH
from devicetest.controllers.device import HarmonyBase
from devicetest.core.constants import DeviceConstants
from devicetest.controllers.device import HarmonyProxy
from devicetest.utils.time_util import TimeHandler
from devicetest.utils.util import get_local_ip_address
from devicetest.utils.util import compare_version
from devicetest.core.exception import AppInstallError
from devicetest.core.exception import RPCException
from devicetest.core.exception import CreateUiDriverFailError
from devicetest.core.exception import ConnectAccessibilityFailError
from devicetest.error import ErrorCategory
from devicetest.error import ErrorMessage
from devicetest.utils.util import compare_versions_by_product
from devicetest.utils.util import check_device_file_md5

DEVICETEST_HAP_PACKAGE_NAME = "com.ohos.devicetest"
DEVICETEST_HOS_HAP_PACKAGE_NAME = "com.ohos.devicetest.plrdtest"
DEVICETEST_HAP_ENTRY_NAME = "entry"
MODIFY_TIME_STR = "Modify: "
BASE_TIME = "2023-10-01 00:00:00"
AGENT_CLEAR_PATH = ["app", "commons-", "agent", "libagent_antry"]


class OpenHarmony(HarmonyBase):
    def __init__(self, port=DeviceConstants.PORT, addr=DeviceConstants.HOST,
                 timeout=300, device=None):
        super(OpenHarmony, self).__init__(port, addr, timeout, device)
        self._device = device
        self.log = device.log
        self.log.debug('init OpenHarmony root')

    def __getattr__(self, name):
        base_proxy = HarmonyProxy(self, name)
        setattr(self, name, base_proxy)
        return base_proxy

    def _rpc(self, module, method, *args):
        data = {'module': module,
                'method': method,
                'params': args,
                'call': "xdevice",
                'request_id': TimeHandler.get_now_datetime()}

        request = json.dumps(data,
                             ensure_ascii=False,
                             separators=(',', ':'))
        return self._safe_send(request, is_bin=False)

    def _abc_rpc(self, module, method, args):
        data = {'module': module,
                'method': method,
                'params': args,
                'call': "xdevice",
                'request_id': TimeHandler.get_now_datetime()}

        request = json.dumps(data,
                             ensure_ascii=False,
                             separators=(',', ':'))
        return self._safe_send(request, is_bin=True)

    def _rpc_for_oh_module(self, module, method, ability_name, *args):
        data = {'module': module,
                'method': method,
                'ability_name': ability_name,
                'params': args,
                'call': "xdevice",
                'request_id': TimeHandler.get_now_datetime()}

        request = json.dumps(data,
                             ensure_ascii=False,
                             separators=(',', ':'))
        return self._safe_send(request, is_bin=False)

    def rpc_for_hypium(self, request):
        new_request = json.loads(request)
        new_request.update({'call': "xdevice"})
        new_request = json.dumps(new_request, ensure_ascii=False, separators=(',', ':'))
        return self._safe_send(new_request,
                               is_bin=self._device.is_bin
                               if hasattr(self._device, "is_bin") else False)

    def _safe_send(self, request, sock=None, isbyte=False, is_bin=False):
        local_host = get_local_ip_address()
        if local_host is None:
            local_host = self.device.host
        send_request = json.loads(request)
        send_request["client"] = \
            local_host if local_host is not None else "127.0.0.1"
        send_request = json.dumps(send_request, ensure_ascii=False, separators=(',', ':'))
        try:
            self.log.debug("sendRequest: {}".format(send_request))
            ret = self._send(send_request, sock=sock, isbyte=isbyte, is_bin=is_bin)
            if b'' == ret:
                return b''
            self.log.debug("@recv rpc msg: {}".format(ret))
            self._check_create_driver(send_request, ret)
        except Exception as e:
            self.log.debug(traceback.format_exc())
            self.log.error("[OpenHarmony] Exception on request to device")
            if self._device is not None:
                if isinstance(e, CreateUiDriverFailError):
                    # 非bin模式下,wukong存不存在都要杀uitest
                    # bin模式下,wukong不存在就报错
                    wukong_exist = self._kill_if_exist_wukong(self._device)
                    if not is_bin:
                        self._kill_uitest(self._device)
                    elif is_bin and not wukong_exist:
                        raise e
                self._wait_accessibility(self._device, request)
                if self._device.reconnecttimes == DeviceConstants.RECONNECT_TIMES:
                    self._device.reconnecttimes = 0
                    error_msg = ErrorMessage.Device.Code_0202001.format("OpenHarmony", DeviceConstants.RECONNECT_TIMES)
                    self.log.error(error_msg)
                    raise RPCException(error_msg) from e

                self.log.debug(
                    "[OpenHarmony] %d times to reconnect rpc socket "
                    "device: %s" % (self._device.reconnecttimes + 1,
                                    self._device.device_id))
                self._device.reconnecttimes += 1
                if self._check_if_request_is_hypium(request, is_bin):
                    self._device.reconnect(proxy=AgentMode.bin)
                else:
                    self._device.reconnect(proxy=AgentMode.hap)
                self.log.debug(
                    "[OpenHarmony] Send request {} again".format(request))
                # bin模式下hypium的接口直接不重发，直接返回, api为driver.create需要重发
                if not self._check_if_request_is_hypium(request, is_bin, check_driver_create=True):
                    self.log.debug("[OpenHarmony] request is not hypium msg or "
                                   "(is hypium msg and api is create driver) need to resend.")
                    ret = self._safe_send(request, sock=sock, isbyte=isbyte, is_bin=is_bin)
                else:
                    self.log.debug("[OpenHarmony] request is hypium msg, no need to resend.")
                    reply = {"exception": "INTERNAL_ERROR Cannot translate frontend object to backend object"}
                    ret = json.dumps(reply, ensure_ascii=False, separators=(',', ':'))
                if 'b' == ret:
                    return b''
                self.log.debug("@recv rpc msg: {}".format(str(ret)))
                self._device.reconnecttimes = 0


        if ret == "false":
            ret = False
        if ret == "true":
            ret = True

        return ret


    def _check_if_request_is_hypium(self, request, is_bin=False, check_driver_create: bool = False):
        if not is_bin:
            return False
        send_request = json.loads(request)
        if send_request.get("method", None) == "callHypiumApi":
            if check_driver_create:
                params = send_request.get("params", {})
                if params.get("api", None) == "Driver.create":
                    return False
            return True
        return False


    @classmethod
    def _check_uitest_version(cls, device, base_version):
        uitest_version = device.execute_shell_command("/system/bin/uitest --version")
        cls.log.debug("uitest version is {}".format(uitest_version))
        return check_uitest_version(uitest_version, base_version)

    @classmethod
    def init_agent_resource(cls, device):
        cls._device = device
        cls.log = device.log
        if cls._device.check_if_bin():
            cls._init_so_resource(device)
        else:
            cls._init_abc_resource(device)

    @classmethod
    def _init_so_resource(cls, device):
        folder_path = os.path.join(RESOURCE_PATH, 'res', 'prototype', 'native')
        file_postfix = ".so"
        device_agent_path = "/data/local/tmp/agent.so"
        cls.log.debug("{}".format("Init native agent..."))
        arch_info = cls._device.execute_shell_command("file /system/bin/uitest")
        cls.log.debug("{}".format(arch_info))
        if "x86_64" in arch_info:
            file_postfix = ".x86_64_so"
        agent_filename = ""
        agent_path = ""
        local_link = "0.0"
        # find local agent file
        for i in os.listdir(folder_path):
            if os.path.splitext(i)[1] == file_postfix:
                agent_path = os.path.join(folder_path, i)
                matcher = re.search(r'\d{1,3}[.]\d{1,3}[.]\d{1,3}', i)
                if matcher:
                    agent_filename = i
                    ver = matcher.group(0)
                    temp1 = tuple(ver.split('.'))
                    temp2 = tuple(local_link.split('.'))
                    if temp2 < temp1:
                        local_link = ver
        # 获取设备端的版本号
        device_ver_info = cls._device.execute_shell_command(
            "cat {} | grep -a UITEST_AGENT_LIBRARY".format(device_agent_path))
        cls.log.debug("{}".format(device_ver_info))
        matcher = re.search(r'\d{1,3}[.]\d{1,3}[.]\d{1,3}', device_ver_info)
        device_link = matcher.group(0) if matcher else "0.0.0"

        need_update = False
        cls.log.debug("local agent version {}, device agent version {}".format(local_link, device_link))
        device_link = tuple(device_link.split("."))
        local_link = tuple(local_link.split("."))
        if device_link < local_link:
            need_update = True
        if need_update:
            cls.log.debug("start update agent, path is {}".format(agent_path))
            # if uitest running kill first
            proc_pids = device.get_devicetest_proc_pid()
            for index, pid in enumerate(proc_pids):
                if pid != "":
                    if index == 2 and device._proxy:
                        continue
                    else:
                        device.execute_shell_command('kill -9 {}'.format(pid))
            # check if has old link file
            for file in AGENT_CLEAR_PATH:
                cls._device.execute_shell_command("rm /data/local/tmp/{}*".format(file))
            cls._device.push_file(agent_path, device_agent_path)
            cls.log.debug("Update agent finish.")
        else:
            cls.log.debug("Device agent is up to date!")

    @classmethod
    def _init_abc_resource(cls, device):
        folder_path = os.path.join(RESOURCE_PATH, 'res', 'prototype', 'abc')
        file_postfix = ".abc"
        device_agent_path = "/data/local/tmp/app.abc"
        cls.log.debug("{}".format("Init abc agent..."))
        agent_filename = ""
        agent_path = ""
        local_link = "0.0"
        # find local agent file
        for i in os.listdir(folder_path):
            if os.path.splitext(i)[1] == file_postfix:
                agent_path = os.path.join(folder_path, i)
                matcher = re.search(r'\d{1,3}[.]\d{1,3}', i)
                if matcher:
                    agent_filename = i
                    ver = matcher.group(0)
                    temp1 = ver.split('.')
                    temp2 = local_link.split('.')
                    for index, _ in enumerate(temp1):
                        if int(temp2[index]) < int(temp1[index]):
                            local_link = ver
                            break
        # check if need update
        device_link_info = cls._device.execute_shell_command("ls -l {}".format(device_agent_path)).strip()
        cls.log.debug("{}".format(device_link_info))
        # 获取版本号
        matcher = re.search(r'\d{1,3}[.]\d{1,3}', device_link_info)
        device_link = matcher.group(0) if matcher else "0.0"
        # 获取delta时间
        link_file_info = device_link_info.split("->")[-1] if "->" in device_link_info else None
        device_delta_time = 0 if link_file_info is None \
            else link_file_info.split("-")[len(link_file_info.split("-")) - 1][:-len(file_postfix)] \
            if len(link_file_info.split("-")) > 2 else 0
        need_update = False
        cls.log.debug("local agent version {}, device agent version {}".format(local_link, device_link))
        device_link = device_link.split(".")
        local_link = local_link.split(".")
        for index, _ in enumerate(device_link):
            if int(device_link[index]) < int(local_link[index]):
                need_update = True
                break
            elif int(device_link[index]) > int(local_link[index]):
                break
        if not need_update and link_file_info:
            # 计算agent文件的修改时间
            delta_time = cls._get_agent_delta_time(link_file_info)
            # 版本号一致还需要判断是否有被修改
            # 若不相等则被修改过
            if delta_time is None or delta_time != int(device_delta_time):
                need_update = True
                cls.log.debug("Agent file has been modified!")
        if need_update:
            cls.log.debug("start update agent, path is {}".format(agent_path))
            # if uitest running kill first
            proc_pids = device.get_devicetest_proc_pid()
            for pid in proc_pids:
                if pid != "":
                    device.execute_shell_command('kill -9 {}'.format(pid))
            # check if has old link file
            for file in AGENT_CLEAR_PATH:
                cls._device.execute_shell_command("rm /data/local/tmp/{}*".format(file))
            cls._device.push_file(agent_path, "/data/local/tmp/{}".format(agent_filename))
            # rename agent file, add modify time to file name
            delta_time = cls._get_agent_delta_time("/data/local/tmp/{}".format(agent_filename))
            if delta_time is None:
                raise FileExistsError("Get agent modify time error!")
            cls._device.execute_shell_command("mv /data/local/tmp/{} /data/local/tmp/{}-{}{}"
                                              .format(agent_filename, agent_filename[:-len(file_postfix)],
                                                      delta_time, file_postfix))
            # link file to /data/local/tmp/agent.so or /data/local/tmp/app.abc
            cls._device.execute_shell_command("ln -s /data/local/tmp/{}-{}{} {}".
                                              format(agent_filename[:-len(file_postfix)], delta_time,
                                                     file_postfix, device_agent_path))
            cls.log.debug("Update agent finish.")
        else:
            cls.log.debug("Device agent is up to date!")
        device.execute_shell_command(
            "param set persist.ace.testmode.enabled 1")

    @classmethod
    def _get_agent_delta_time(cls, agent_file_path):
        delta_time = None
        base_time_array = time.strptime(BASE_TIME, "%Y-%m-%d %H:%M:%S")
        base_time_stamp = time.mktime(base_time_array)
        # 计算agent文件的修改时间
        stat_info = cls._device.execute_shell_command(
            "stat {}".format(agent_file_path)).strip().split("\n")
        modify_time = None
        for info in stat_info:
            if info.startswith(MODIFY_TIME_STR):
                index = info.find(".")
                modify_time = info[len(MODIFY_TIME_STR):index]
                break
        if modify_time:
            modify_time_array = time.strptime(modify_time, "%Y-%m-%d %H:%M:%S")
            modify_time_stamp = time.mktime(modify_time_array)
            delta_time = (int)(modify_time_stamp - base_time_stamp)
        return delta_time


    @classmethod
    def _kill_if_exist_wukong(cls, device) -> bool:
        # 获取wukong程号然后杀掉,不存在就返回False
        pids_res = device.execute_shell_command("pidof wukong")
        cls.log.debug("wukong pid : {}".format(pids_res))
        if not pids_res:
            return False
        for pid in pids_res.strip().split("\n"):
            device.execute_shell_command("kill -9 {}".format(pid))
        time.sleep(1)
        return True

    @classmethod
    def _kill_uitest(cls, device):
        # 获取uitest程号然后杀掉
        pids_res = device.execute_shell_command("ps -A | grep uitest")
        cls.log.debug("uitest pid : {}".format(pids_res))
        if not pids_res:
            return
        for process in pids_res.strip().split("\n"):
            m = re.match("^(\\s*\\d+\\s+)", process)
            if m:
                pid = m.group(1).strip()
                if pid.isdigit():
                    device.execute_shell_command("kill -9 {}".format(pid))
        time.sleep(1)

    @classmethod
    def _wait_accessibility(cls, device, request: str):
        # 非UI操作不等待
        if not hasattr(device, "is_oh") or not cls._check_ui_action(request):
            return
        # uitest版本大于5.0.0.0不需要等待无障碍
        uitest_version = device.execute_shell_command("/system/bin/uitest --version")
        base_version = tuple("5.0.0.0".split("."))
        if check_uitest_version(uitest_version, base_version):
            return
        # 产品版本大于5.0.0.22不需要等待无障碍
        product_info = device.execute_shell_command("param get const.product.software.version").strip().split(" ")
        if len(product_info) == 2:
            device_version = product_info[1]
            if compare_versions_by_product(device_version, "5.0.0.22"):
                return
        start_time = time.time()  # 获取当前时间
        is_accessibility_alive = False
        while time.time() - start_time < 60:  # 60秒内持续等待无障碍
            # 检测无障碍是否已被拉起
            pids_res = device.execute_shell_command("pidof accessibility")
            cls.log.debug("waiting accessibility")
            if pids_res:
                is_accessibility_alive = True
                cls.log.debug("accessibility pid is {}".format(pids_res))
                break
            time.sleep(2)
        if not is_accessibility_alive:
            raise ConnectAccessibilityFailError(ErrorMessage.Device.Code_0202006)

    @classmethod
    def _check_create_driver(cls, request, ret):
        try:
            send_request_dict = json.loads(request)
            params = send_request_dict.get("params")
            if not params or not isinstance(ret, str):
                return
            datas_dict = json.loads(ret)
            if not isinstance(datas_dict, dict):
                return
        except Exception as error:
            cls.log.debug("parse json error: {}".format(error))
            return
        exception = datas_dict.get("exception")
        if exception and "INTERNAL_ERROR Cannot find backend method" in exception:
            raise CreateUiDriverFailError(ErrorMessage.Device.Code_0202007)
        if exception and "RET_ERR_CONNECTION_EXIST" in exception:
            raise CreateUiDriverFailError(ErrorMessage.Device.Code_0202009)
        for param in params:
            if not param:
                continue
            if "Driver.create" in param:
                if not datas_dict.get("result"):
                    error_msg = ErrorMessage.Device.Code_0202008
                    cls.log.error(error_msg)
                    raise CreateUiDriverFailError(error_msg)

    @classmethod
    def _check_ui_action(cls, request: str):
        # 判断当前请求是否为UI操作
        try:
            if not isinstance(request, str):
                return False
            send_request_dict = json.loads(request)
            module = send_request_dict.get("module")
            if not module:
                return False
        except Exception as error:
            cls.log.debug("parse json error: {}".format(error))
            return False
        return module == "com.ohos.devicetest.hypiumApiHelper" or module == "com.ohos.devicetest.UiTestDeamon"
